{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dynamic Deep Hit on PBC Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The longitudinal PBC dataset comes from the Mayo Clinic trial in primary biliary cirrhosis (PBC) of the liver conducted between 1974 and 1984 (Refer to https://stat.ethz.ch/R-manual/R-devel/library/survival/html/pbc.html)\n",
    "\n",
    "In this notebook, we will apply Dynamic Deep Hit for survival prediction on the PBC data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('..')\n",
    "sys.path.append('../DeepSurvivalMachines')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load the PBC Dataset\n",
    "\n",
    "The package includes helper functions to load the dataset.\n",
    "\n",
    "X represents an np.array of features (covariates),\n",
    "T is the event/censoring times and,\n",
    "E is the censoring indicator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dsm import datasets\n",
    "x, t, e = datasets.load_dataset('PBC', sequential = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Additionally, the DeepHit model aligns the prediction given time 0, it is therefore necessary to have the time of last observation. We choose to have the last observation as time 0 and give an array of same size than t with 0."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compute horizons at which we evaluate the performance of Dynamic Deep Hit\n",
    "\n",
    "Survival predictions are issued at certain time horizons. Here we will evaluate the performance\n",
    "of DDH to issue predictions at the 25th, 50th and 75th event time quantile as is standard practice in Survival Analysis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "horizons = [0.25, 0.5, 0.75]\n",
    "times = np.quantile([t_[-1] for t_, e_ in zip(t, e) if e_[-1] == 1], horizons).tolist()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Splitting the data into train, test and validation sets\n",
    "\n",
    "We will train RDSM on 70% of the Data, use a Validation set of 10% for Model Selection and report performance on the remaining 20% held out test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = len(x)\n",
    "\n",
    "tr_size = int(n*0.70)\n",
    "vl_size = int(n*0.10)\n",
    "te_size = int(n*0.20)\n",
    "\n",
    "x_train, x_test, x_val = np.array(x[:tr_size], dtype = object), np.array(x[-te_size:], dtype = object), np.array(x[tr_size:tr_size+vl_size], dtype = object)\n",
    "t_train, t_test, t_val = np.array(t[:tr_size], dtype = object), np.array(t[-te_size:], dtype = object), np.array(t[tr_size:tr_size+vl_size], dtype = object)\n",
    "e_train, e_test, e_val = np.array(e[:tr_size], dtype = object), np.array(e[-te_size:], dtype = object), np.array(e[tr_size:tr_size+vl_size], dtype = object)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setting the parameter grid\n",
    "\n",
    "Lets set up the parameter grid to tune hyper-parameters. We will tune the number of underlying survival distributions, \n",
    "($K$), the distribution choices (Log-Normal or Weibull), the learning rate for the Adam optimizer between $1\\times10^{-3}$ and $1\\times10^{-4}$, the number of hidden nodes per layer $50, 100$ and $2$, the number of layers $3, 2$ and $1$ and the type of recurrent cell (LSTM, GRU, RNN)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import ParameterGrid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "layers = [[], [100], [100, 100], [100, 100, 100]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "param_grid = {\n",
    "              'layers_rnn': [2, 3],\n",
    "              'hidden_long': layers,\n",
    "              'hidden_rnn': [50, 100],\n",
    "              'hidden_att': layers,\n",
    "              'hidden_cs': layers,\n",
    "              'sigma': [0.1, 1, 3],\n",
    "              'learning_rate' : [1e-3],\n",
    "             }\n",
    "params = ParameterGrid(param_grid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model Training and Selection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from ddh.ddh_api import DynamicDeepHit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "models = []\n",
    "for param in params:\n",
    "    model = DynamicDeepHit(\n",
    "                layers_rnn = param['layers_rnn'],\n",
    "                hidden_rnn = param['hidden_rnn'], \n",
    "                long_param = {'layers': param['hidden_long'], 'dropout': 0.3}, \n",
    "                att_param = {'layers': param['hidden_att'], 'dropout': 0.3}, \n",
    "                cs_param = {'layers': param['hidden_cs'], 'dropout': 0.3},\n",
    "                sigma = param['sigma'],\n",
    "                split = [0] + times + [np.max([t_.max() for t_ in t])])\n",
    "    # The fit method is called to train the model\n",
    "    model.fit(x_train, t_train, e_train, iters = 10, \n",
    "              learning_rate = param['learning_rate'])\n",
    "    models.append([[model.compute_nll(x_val, t_val, e_val), model]])\n",
    "best_model = min(models)\n",
    "model = best_model[0][1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_model = min(models)\n",
    "model = best_model[0][1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out_risk = model.predict_risk(x_test, times, all_step = True)\n",
    "out_survival = model.predict_survival(x_test, times, all_step = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluation\n",
    "\n",
    "We evaluate the performance of RDSM in its discriminative ability (Time Dependent Concordance Index and Cumulative Dynamic AUC) as well as Brier Score on the concatenated temporal data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sksurv.metrics import concordance_index_ipcw, brier_score, cumulative_dynamic_auc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cis = []\n",
    "brs = []\n",
    "\n",
    "et_train = np.array([(e_train[i][j], t_train[i][j]) for i in range(len(e_train)) for j in range(len(e_train[i]))],\n",
    "                 dtype = [('e', bool), ('t', float)])\n",
    "et_test = np.array([(e_test[i][j], t_test[i][j]) for i in range(len(e_test)) for j in range(len(e_test[i]))],\n",
    "                 dtype = [('e', bool), ('t', float)])\n",
    "et_val = np.array([(e_val[i][j], t_val[i][j]) for i in range(len(e_val)) for j in range(len(e_val[i]))],\n",
    "                 dtype = [('e', bool), ('t', float)])\n",
    "\n",
    "for i, _ in enumerate(times):\n",
    "    cis.append(concordance_index_ipcw(et_train, et_test, out_risk[:, i], times[i])[0])\n",
    "brs.append(brier_score(et_train, et_test, out_survival, times)[1])\n",
    "roc_auc = []\n",
    "for i, _ in enumerate(times):\n",
    "    roc_auc.append(cumulative_dynamic_auc(et_train, et_test, out_risk[:, i], times[i])[0])\n",
    "for horizon in enumerate(horizons):\n",
    "    print(f\"For {horizon[1]} quantile,\")\n",
    "    print(\"TD Concordance Index:\", cis[horizon[0]])\n",
    "    print(\"Brier Score:\", brs[0][horizon[0]])\n",
    "    print(\"ROC AUC \", roc_auc[horizon[0]][0], \"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
